package com.taobao.android.mnndemo;

import com.taobao.android.mnn.MNNNetInstance;
import com.taobao.android.mnn.MNNImageProcess;

public class MNNInference {
    private MNNNetInstance netInstance;
    private MNNNetInstance.Session session;

    public MNNInference(String modelPath) {
        try {
            // 加载 MNN 模型
            netInstance = MNNNetInstance.createFromFile(modelPath);
            // 创建会话
            session = netInstance.createSession(null);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public float[][][] infer(float[] inputData) {
        if (session == null) {
            return null;
        }

        // 获取输入 tensor
        MNNNetInstance.Session.Tensor inputTensor = session.getInput(null);

        // 填充输入数据
        inputTensor.setInputFloatData(inputData);

        // 运行推理
        session.run();

        // 获取输出 tensor
        MNNNetInstance.Session.Tensor outputTensor = session.getOutput(null);

        // 获取输出数据并解析
        float[] outputData = outputTensor.getFloatData();

        // 解析输出
        float[][][] keypoints = new float[1][17][3];
        for (int i = 0; i < 17; i++) {
            keypoints[0][i][0] = outputData[i * 3];     // x 坐标
            keypoints[0][i][1] = outputData[i * 3 + 1]; // y 坐标
            keypoints[0][i][2] = outputData[i * 3 + 2]; // 置信度
        }

        return keypoints;
    }

    public void release() {
        if (netInstance != null) {
            netInstance.release();
        }
    }
}